import numpy as np
import matplotlib.pyplot as plt
import menu_cost as mc
import scipy.optimize as opt
from simple_block import simple, exog
import jacobian as jac


def phillips_curve(J):
    return exog(J)


@simple
def model(y, n, ns, w, i, pi, freq, delta, sig, phi, phi_pi, phi_y, N, Ns, xi, eps_i):
    euler_res = y - (y(1) - (1 / sig) * (i - pi(1)))
    labor_res = phi * ns + sig * y - w
    prod_res = y - delta - n
    mkt_res = ns - (N / Ns) * n - (xi / Ns) * freq
    taylor_res = i - (phi_pi * pi + phi_y * y + eps_i)
    return euler_res, labor_res, prod_res, mkt_res, taylor_res


@simple
def pc_inputs(y, w, sig, elast):
    discount = -sig * y + w
    shift = -(elast * w) + y
    gap = w
    return discount, gap, shift


def compute_ss(parameters, ss, sig=1, phi=5, Y=1):
    elast = parameters['elast']
    la = parameters['la']
    Delta = ss['Delta']
    f = ss['stats']['freq']

    def ss_res(w, b):
        shift = np.log(Y * (w ** (-elast)))

        Pi = Y * (w ** (1 - elast)) * ss['Profits']
        xi = np.exp(parameters['mu_k'] + shift)

        N = Y / Delta
        Ns = (w * Y ** (-sig) / b) ** (1 / phi)
        Nd = N + xi * (f - la)

        res = np.array([Nd - Ns, Y - w * N - Pi])

        return res, shift, Delta, f, la, xi, N, Pi, Nd, Ns, b

    x0 = np.array([0.75, 0.7])
    fobj = lambda x: ss_res(x[0], x[1])[0]
    result = opt.fsolve(fobj, x0)
    w, b = result

    _, shift, Delta, f, la, xi, N, Pi, Nd, Ns, b = ss_res(w, b)

    ss_nk = {'y': 0, 'n': 0, 'ns': 0, 'w': 0, 'i': 0, 'pi': 0, 'freq': 0, 'delta': 0,
             'sig': sig, 'phi': phi, 'phi_pi': 1.5, 'phi_y': 0, 'N': N, 'Ns': Ns, 'xi': xi, 'elast': elast,
             'eps_i': 0}

    return ss_nk

# block_list = [model, pc_inputs, exog(J)]
unknowns = ['y', 'n', 'ns', 'w', 'i']
targets = ['euler_res', 'labor_res', 'prod_res', 'mkt_res', 'taylor_res']
